# -*- coding:utf8 -*-

import os
import random
# from PIL import Image

allowed_names = ['jpg', 'JPG']  # 需要随机替换的后缀名列表
COUNT_NUMBER = 0
TRAIN_NUMBER_RATE = 0.9
SAVE_PATH = '/home/user/my/'
INPUT_PATH = '/usr/VOCdevkit/VOC2019/Annotation'
TRAIN_TXT_NAME = 'train.txt'
VAL_TXT_NAME = 'val.txt'

def file_name(name):
    a = name.split('.', 1)[0]
    return name.split('.', 1)[0]

def file_constract(path):
    files = os.listdir(path)  # 获取当前目录的所有文件及文件夹
    file_number = len(files)
    train_number = int(file_number * TRAIN_NUMBER_RATE)
    train_files = random.sample(files, train_number)

    train_file_names = [file_name(train_file) for train_file in train_files]
    # for train_file in train_files:
    #     b = file_name(train_file)
    #     c = 1
    # files.remove(train_file for train_file in train_files)
    for train_file in train_files:
        files.remove(train_file)
    val_file_names = [file_name(val_file) for val_file in files]

    return train_file_names, val_file_names

if __name__ == '__main__':
    train_names, val_names = file_constract(INPUT_PATH)
    with open(SAVE_PATH+TRAIN_TXT_NAME, 'a') as f:
        for train_name in train_names:
            f.write(str(train_name)+'\n')
        f.close()
    with open(SAVE_PATH+VAL_TXT_NAME, 'a') as f:
        for val_name in val_names:
            f.write(str(val_name)+'\n')
        f.close()